import torch
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")
sys.path.append(r"/home/cyw/projects/function_sim_project/function_sim")
import shelve
import torch 
from easySample import easySample
import numpy as np

class sample_function_embedding():
    """
        获得样本经过模型训练后的各个函数的嵌入值(新增操作每一层)
        input:模型，样本a，样本b
        output:两个样本的图节点的嵌入值
    """
    def __init__(self,model,namea,nameb) -> None:
        """
            namea为样本的MD5值
        """
        # shelve 名字区分大小写,大写的时候生成了一个新的shelve文件
        namea=namea.lower()
        nameb=nameb.lower()

        self.model=torch.load(model)
        sample=easySample()
        
        sample_a=sample.get_sample(namea,"functionSim")
        sample_b=sample.get_sample(nameb,"functionSim")
        
        x_adj=torch.tensor([sample_a["adj"]],dtype=torch.float64).cuda()
        x_att=torch.tensor([sample_a["att"]],dtype=torch.float64).cuda()

        x_type=torch.tensor([sample_a["vtype"]],dtype=torch.float64).cuda()
        y_adj=torch.tensor([sample_b["adj"]],dtype=torch.float64).cuda()
        y_att=torch.tensor([sample_b["att"]],dtype=torch.float64).cuda()
        y_type=torch.tensor([sample_b["vtype"]],dtype=torch.float64).cuda()
        score1=self.model(x_adj,x_att,x_type,y_adj,y_att,y_type)
        print(score1)



if __name__ == "__main__":
    # 需要开启保存embedding，使用后请关闭，否则磁盘会被塞满
    # !!!!!!!!切换的时候配置信息中的设置也需要改变！！！！！！！！！！  难得写代码了，用的时候手动注释一下！！！！！！
    functionSimModel=[r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge_best_4_8.pth",
     r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_with_hete_model.pth",
     r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_zero_model.pth"]
    
    aList=[["eec55f1295d0cce40bbba7be231fdddb","6abf24e3aa3e971e6461bcfea35fe950"],
           ["e940f6319ad362ec49af956cafaf8752","d72fe801683ee70dccaba657984b7997"],
           ["1da46f95579e467013e94ca4a5117de0","496430ce3d8623147be33937f92265a7"],
           ["7d9d0d28eaed286c029947e585993b26","1da46f95579e467013e94ca4a5117de0"],
           ["cfc34d91c9114483f3101c5049afb10d","e8a683916a98d7f132fd8d05d52ea66c"],

           ["74ec6006ed76a9bed14c140e53c7efd4","bb6f9c2aa88245746e24d7672016afb7"],
           ["b68a07cf3f63172901ce8298e78d7a74","63ba29be215b532731392511dcd4363e"],
           ["16f8a4cdd845bc4680c2822d8dbc10c4","bb6f9c2aa88245746e24d7672016afb7"],
           ["698ff16b0e824db5874acee0f7324bb2","a6bd7c5a22dbf503e7d0510b02494a95"],
           ["ec45bbfb38da5b65a58eb018b0122285","d160839802998ccabdb7ca066ae2eb3b"],

           ["50c64842dc9e99f259b4a247762b918e","bb6f9c2aa88245746e24d7672016afb7"],
           ["00a109d00d01fc2be04738ee9dcc65c2","22c1b1924703ad7d0e142c5c8ebc57e0"],
           ["81cbb6f4b1cc8d5e2167d106657de095","b33da0446151ac45b2e07491adde796d"],
           ["ce1a3ebdde8eb6d24e2c91bd3ce1da72","c1eb31c7c0ff87c5b182f8e3a447028a"],
           ["2db1afa1d21b8bbe83d7208f4c9942b3","e3fb5232f88b4a62b0d80d31e4828770"],

           ["f846e2ca4c4a36a00dbc28abd892e4b9","cbb2da21894d34cdbadaa26d1ec03c01"]
           ]
    ind = 0
    model=functionSimModel[2]



    sample_a=aList[ind][0]
    sample_b=aList[ind][1]
    a=sample_function_embedding(model,sample_a,sample_b)
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
    
    
